import os
import sys
import subprocess
import numpy as np
import pandas as pd
from sklearn.preprocessing import MinMaxScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split
import joblib
from loguru import logger

# 配置 loguru
logger.remove()  # 移除默认的处理器
logger.add(sys.stdout, format="{time:YYYY-MM-DD HH:mm:ss.SSS} | {level:<8} | {name}:{function}:{line} - {message}", level="INFO")

# ---------------- 配置 ----------------
current_dir = os.path.dirname(os.path.abspath(__file__))
DATA_FILE = os.path.join(current_dir, "fc3d_history.csv")
MODEL_PATH = os.path.join(current_dir, "random_rf_3d_model.pkl")
SCALER_PATH = os.path.join(current_dir, "random_scaler_X.pkl")
WINDOW_SIZE = 10

# 获取项目根目录
project_root = os.path.abspath(os.path.join(current_dir, '..', '..'))
sys.path.append(project_root)

def fetch_data_if_not_exists():
    """
    检查 CSV 文件是否存在，如果不存在，则调用 fetch_3d_data.py 获取数据
    """
    if not os.path.exists(DATA_FILE):
        logger.info(f"数据文件 {DATA_FILE} 不存在，开始获取数据...")
        fetch_script = os.path.join(current_dir, 'fetch_3d_data.py')
        if not os.path.exists(fetch_script):
            logger.error(f"数据获取脚本不存在: {fetch_script}")
            sys.exit(1)
        try:
            # 使用当前运行的 Python 解释器
            python_executable = sys.executable
            logger.info(f"运行数据获取脚本: {fetch_script} 使用解释器: {python_executable}")
            subprocess.run([python_executable, fetch_script], check=True)
            logger.info("数据获取完成。")
        except subprocess.CalledProcessError as e:
            logger.error(f"运行数据获取脚本失败: {e}")
            sys.exit(1)
    else:
        logger.info(f"数据文件 {DATA_FILE} 已存在。")

def preprocess_data(data, window_size):
    """
    预处理3D数据
    3D彩票有3个数字位置，每个位置的数字范围是0-9
    """
    features, labels = [], []
    
    # 确保数据列存在
    required_columns = ['num_1', 'num_2', 'num_3']
    for col in required_columns:
        if col not in data.columns:
            raise ValueError(f"数据格式错误：缺少列 {col}")
    
    # 获取号码数据
    number_data = data[required_columns].values
    
    for i in range(len(number_data) - window_size):
        # 特征：选取窗口内的3个数字
        feature_window = number_data[i:i + window_size].flatten()  # 展平为一维数组
        features.append(feature_window)

        # 标签：下一期的3个数字
        next_numbers = number_data[i + window_size]
        labels.append(next_numbers)

    # 转换为 NumPy 数组并进行缩放
    features_np = np.array(features)  # 形状: (num_samples, window_size * 3)
    scaler_X = MinMaxScaler()
    features_scaled = scaler_X.fit_transform(features_np)

    labels_np = np.array(labels)  # 形状: (num_samples, 3)

    return features_scaled, labels_np, scaler_X

def train_model():
    fetch_data_if_not_exists()

    if not os.path.exists(DATA_FILE):
        logger.error(f"数据文件不存在: {DATA_FILE}")
        sys.exit(1)

    # 数据加载
    logger.info("加载3D数据...")
    data = pd.read_csv(DATA_FILE)

    # 数据预处理
    features, labels, scaler_X = preprocess_data(data, WINDOW_SIZE)

    # 打印特征和标签的形状
    logger.info(f"特征形状: {features.shape}")
    logger.info(f"标签形状: {labels.shape}")

    # 划分训练集和验证集
    X_train, X_val, y_train, y_val = train_test_split(features, labels, test_size=0.1, random_state=42)

    # 模型初始化 - 为3个数字位置分别创建模型
    models = []
    for i in range(3):  # 3个数字位置
        model = RandomForestClassifier(n_estimators=100, random_state=42)
        models.append(model)

    # 训练过程
    logger.info("开始模型训练...")
    
    accuracies = []
    for i in range(3):  # 分别训练每个位置的模型
        # 训练模型
        models[i].fit(X_train, y_train[:, i])
        preds = models[i].predict(X_val)
        
        # 计算准确率
        accuracy = accuracy_score(y_val[:, i], preds)
        accuracies.append(accuracy)
        logger.info(f"位置 {i+1} 模型训练完成，验证集准确率: {accuracy:.4f}")

    # 保存模型和缩放器
    os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True)
    joblib.dump({
        "models": models,
        "positions": 3
    }, MODEL_PATH)
    joblib.dump(scaler_X, SCALER_PATH)
    logger.info(f"模型已保存到 {MODEL_PATH}")
    logger.info(f"缩放器已保存到 {SCALER_PATH}")
    
    # 输出平均准确率
    avg_accuracy = np.mean(accuracies)
    logger.info(f"平均准确率: {avg_accuracy:.4f}")

def load_model_and_scaler(model_path, scaler_path):
    """ 加载模型和缩放器 """
    model = joblib.load(model_path)
    scaler_X = joblib.load(scaler_path)
    return model, scaler_X

def prepare_new_data(data, window_size):
    """ 准备新数据用于预测 """
    if len(data) < window_size:
        logger.error(f"数据行数不足 {window_size} 行。当前数据行数: {len(data)}")
        return None
    
    # 确保数据列存在
    required_columns = ['num_1', 'num_2', 'num_3']
    for col in required_columns:
        if col not in data.columns:
            raise ValueError(f"数据格式错误：缺少列 {col}")
    
    # 获取号码数据
    number_data = data[required_columns].values
    
    features = []
    for i in range(len(number_data) - window_size + 1):
        # 特征：选取窗口内的3个数字
        feature_window = number_data[i:i + window_size].flatten()
        features.append(feature_window)
    
    # 检查 features 列表是否为空
    if not features:
        logger.error("Features list is empty. Ensure there are enough data rows for the window size.")
        return None
    
    features_np = np.array(features)  # 形状: (num_samples, window_size * 3)
    return features_np

def predict_next_draw(model, scaler_X, data, window_size):
    """ 使用模型预测下一期的3D彩票号码 """
    features_scaled = prepare_new_data(data, window_size)
    if features_scaled is None:
        return None
    
    # 打印缩放后的特征数据
    logger.debug(f"特征缩放形状: {features_scaled.shape}")
    
    models = model["models"]
    
    # 预测每个位置的数字
    predictions = []
    for i in range(len(models)):  # 对每个位置进行预测
        preds = models[i].predict(features_scaled)
        predictions.append(preds)
    
    # 转置结果，使每个样本的预测结果为一行
    predictions = np.array(predictions).T
    return predictions

if __name__ == "__main__":
    logger.info("开始训练3D随机森林模型...")
    train_model()
    logger.info("3D随机森林模型训练完成。")
    
    # 示例预测代码（可选）
    # logger.info("加载数据...")
    # data = pd.read_csv(DATA_FILE)
    # 
    # logger.info("加载模型和缩放器...")
    # model, scaler_X = load_model_and_scaler(MODEL_PATH, SCALER_PATH)
    #
    # logger.info("开始预测下一期的3D彩票号码...")
    # predictions = predict_next_draw(model, scaler_X, data, WINDOW_SIZE)
    # if predictions is not None:
    #     logger.info(f"预测结果: {predictions[-1]}")  # 显示最后一期的预测
    # else:
    #     logger.error("预测失败。")